// David Eberly, Geometric Tools, Redmond WA 98052
// Copyright (c) 1998-2019
// Distributed under the Boost Software License, Version 1.0.
// http://www.boost.org/LICENSE_1_0.txt
// http://www.geometrictools.com/License/Boost/LICENSE_1_0.txt
// File Version: 3.0.2 (2019/04/12)

#pragma once

#include <array>
#include <cmath>
#include <cstdint>

// An implementation of the QR algorithm described in "Matrix Computations,
// 2nd edition" by G. H. Golub and C. F. Van Loan, The Johns Hopkins
// University Press, Baltimore MD, Fourth Printing 1993.  In particular,
// the implementation is based on Chapter 7 (The Unsymmetric Eigenvalue
// Problem), Section 7.5 (The Practical QR Algorithm).  The algorithm is
// specialized for the companion matrix associated with a cubic polynomial.

namespace gte
{

template <typename Real>
class CubicRootsQR
{
public:
    typedef std::array<std::array<Real, 3>, 3> Matrix;

    // Solve p(x) = c0 + c1 * x + c2 * x^2 + x^3 = 0.
    uint32_t operator() (uint32_t maxIterations, Real c0, Real c1, Real c2,
        uint32_t& numRoots, std::array<Real, 3>& roots);

    // Compute the real eigenvalues of the upper Hessenberg matrix A.  The
    // matrix is modified by in-place operations, so if you need to remember
    // A, you must make your own copy before calling this function.
    uint32_t operator() (uint32_t maxIterations, Matrix& A,
        uint32_t& numRoots, std::array<Real, 3>& roots);

private:
    void DoIteration(std::array<Real, 3> const& V, Matrix& A);

    template <int N>
    std::array<Real, N> House(std::array<Real, N> const& X);

    template <int N>
    void RowHouse(int rmin, int rmax, int cmin, int cmax,
        std::array<Real, N> const& V, std::array<Real, N> const& MV, Matrix& A);

    template <int N>
    void ColHouse(int rmin, int rmax, int cmin, int cmax,
        std::array<Real, N> const& V, std::array<Real, N> const& MV, Matrix& A);

    void GetQuadraticRoots(int i0, int i1, Matrix const & A,
        uint32_t& numRoots, std::array<Real, 3>& roots);
};


template <typename Real>
uint32_t CubicRootsQR<Real>::operator() (uint32_t maxIterations, Real c0, Real c1, Real c2,
    uint32_t& numRoots, std::array<Real, 3>& roots)
{
    // Create the companion matrix for the polynomial.  The matrix is in upper
    // Hessenberg form.
    Matrix A;
    A[0][0] = (Real)0;
    A[0][1] = (Real)0;
    A[0][2] = -c0;
    A[1][0] = (Real)1;
    A[1][1] = (Real)0;
    A[1][2] = -c1;
    A[2][0] = (Real)0;
    A[2][1] = (Real)1;
    A[2][2] = -c2;

    // Avoid the QR-cycle when c1 = c2 = 0 and avoid the slow convergence
    // when c1 and c2 are nearly zero.
    std::array<Real, 3> V{
        (Real)1,
        (Real)0.36602540378443865,
        (Real)0.36602540378443865 };
    DoIteration(V, A);

    return operator()(maxIterations, A, numRoots, roots);
}

template <typename Real>
uint32_t CubicRootsQR<Real>::operator() (uint32_t maxIterations, Matrix& A,
    uint32_t& numRoots, std::array<Real, 3>& roots)
{
    numRoots = 0;
    std::fill(roots.begin(), roots.end(), (Real)0);

    for (uint32_t numIterations = 0; numIterations < maxIterations; ++numIterations)
    {
        // Apply a Francis QR iteration.
        Real tr = A[1][1] + A[2][2];
        Real det = A[1][1] * A[2][2] - A[1][2] * A[2][1];
        std::array<Real, 3> X{
            A[0][0] * A[0][0] + A[0][1] * A[1][0] - tr * A[0][0] + det,
            A[1][0] * (A[0][0] + A[1][1] - tr),
            A[1][0] * A[2][1] };
        std::array<Real, 3> V = House<3>(X);
        DoIteration(V, A);

        // Test for uncoupling of A.
        Real tr01 = A[0][0] + A[1][1];
        if (tr01 + A[1][0] == tr01)
        {
            numRoots = 1;
            roots[0] = A[0][0];
            GetQuadraticRoots(1, 2, A, numRoots, roots);
            return numIterations;
        }

        Real tr12 = A[1][1] + A[2][2];
        if (tr12 + A[2][1] == tr12)
        {
            numRoots = 1;
            roots[0] = A[2][2];
            GetQuadraticRoots(0, 1, A, numRoots, roots);
            return numIterations;
        }
    }
    return maxIterations;
}

template <typename Real>
void CubicRootsQR<Real>::DoIteration(std::array<Real, 3> const& V, Matrix& A)
{
    Real multV = ((Real)-2) / (V[0] * V[0] + V[1] * V[1] + V[2] * V[2]);
    std::array<Real, 3> MV{ multV * V[0], multV * V[1], multV * V[2] };
    RowHouse<3>(0, 2, 0, 2, V, MV, A);
    ColHouse<3>(0, 2, 0, 2, V, MV, A);

    std::array<Real, 2> Y{ A[1][0], A[2][0] };
    std::array<Real, 2> W = House<2>(Y);
    Real multW = ((Real)-2) / (W[0] * W[0] + W[1] * W[1]);
    std::array<Real, 2> MW{ multW * W[0], multW * W[1] };
    RowHouse<2>(1, 2, 0, 2, W, MW, A);
    ColHouse<2>(0, 2, 1, 2, W, MW, A);
}

template <typename Real>
template <int N>
std::array<Real, N> CubicRootsQR<Real>::House(std::array<Real, N> const & X)
{
    std::array<Real, N> V;
    Real length = (Real)0;
    for (int i = 0; i < N; ++i)
    {
        length += X[i] * X[i];
    }
    length = std::sqrt(length);
    if (length != (Real)0)
    {
        Real sign = (X[0] >= (Real)0 ? (Real)1 : (Real)-1);
        Real denom = X[0] + sign * length;
        for (int i = 1; i < N; ++i)
        {
            V[i] = X[i] / denom;
        }
    }
    else
    {
        V.fill((Real)0);
    }
    V[0] = (Real)1;
    return V;
}

template <typename Real>
template <int N>
void CubicRootsQR<Real>::RowHouse(int rmin, int rmax, int cmin, int cmax,
    std::array<Real, N> const& V, std::array<Real, N> const& MV, Matrix& A)
{
    std::array<Real, 3> W;  // only elements cmin through cmax are used

    for (int c = cmin; c <= cmax; ++c)
    {
        W[c] = (Real)0;
        for (int r = rmin, k = 0; r <= rmax; ++r, ++k)
        {
            W[c] += V[k] * A[r][c];
        }
    }

    for (int r = rmin, k = 0; r <= rmax; ++r, ++k)
    {
        for (int c = cmin; c <= cmax; ++c)
        {
            A[r][c] += MV[k] * W[c];
        }
    }
}

template <typename Real>
template <int N>
void CubicRootsQR<Real>::ColHouse(int rmin, int rmax, int cmin, int cmax,
    std::array<Real, N> const& V, std::array<Real, N> const& MV, Matrix& A)
{
    std::array<Real, 3> W;  // only elements rmin through rmax are used

    for (int r = rmin; r <= rmax; ++r)
    {
        W[r] = (Real)0;
        for (int c = cmin, k = 0; c <= cmax; ++c, ++k)
        {
            W[r] += V[k] * A[r][c];
        }
    }

    for (int r = rmin; r <= rmax; ++r)
    {
        for (int c = cmin, k = 0; c <= cmax; ++c, ++k)
        {
            A[r][c] += W[r] * MV[k];
        }
    }
}

template <typename Real>
void CubicRootsQR<Real>::GetQuadraticRoots(int i0, int i1, Matrix const& A,
    uint32_t& numRoots, std::array<Real, 3>& roots)
{
    // Solve x^2 - t * x + d = 0, where t is the trace and d is the
    // determinant of the 2x2 matrix defined by indices i0 and i1.  The
    // discriminant is D = (t/2)^2 - d.  When D >= 0, the roots are real
    // values t/2 - sqrt(D) and t/2 + sqrt(D).  To avoid potential numerical
    // issues with subtractive cancellation, the roots are computed as
    //   r0 = t/2 + sign(t/2)*sqrt(D), r1 = trace - r0.
    Real trace = A[i0][i0] + A[i1][i1];
    Real halfTrace = trace * (Real)0.5;
    Real determinant = A[i0][i0] * A[i1][i1] - A[i0][i1] * A[i1][i0];
    Real discriminant = halfTrace * halfTrace - determinant;
    if (discriminant >= (Real)0)
    {
        Real sign = (trace >= (Real)0 ? (Real)1 : (Real)-1);
        Real root = halfTrace + sign * std::sqrt(discriminant);
        roots[numRoots++] = root;
        roots[numRoots++] = trace - root;
    }
}

}
